from icecream import ic
import enum
from kor import create_extraction_chain, Object, Text, Number
from typing import List
from kor import from_pydantic
from pydantic import BaseModel, Field
from typing import Optional
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model='gpt-4-turbo-preview',
				 api_key="sk-oYQsYEo1JPME6JEYo3AT9xBupdO3Nqb64sHUNZygh4txMLeC",
				 base_url="https://www.DMXapi.com/v1/")

class Action(enum.Enum):
    play = "play"
    stop = "stop"
    previous = "previous"
    next_ = "next"


class MusicRequest(BaseModel):
    song: Optional[List[str]] = Field(
        default=None, description="The song(s) that the user would like to be played."
    )
    album: Optional[List[str]] = Field(
        default=None, description="The album(s) that the user would like to be played."
    )
    artist: Optional[List[str]] = Field(
        default=None,
        description="The artist(s) whose music the user would like to hear.",
        examples=[("Songs by paul simon", "paul simon")],
    )
    action: Optional[Action] = Field(
        default=None,
        description="The action that should be taken; one of `play`, `stop`, `next`, `previous`",
        examples=[
            ("Please stop the music", "stop"),
            ("play something", "play"),
            ("play a song", "play"),
            ("next song", "next"),
        ],
    )

schema, validator = from_pydantic(MusicRequest)
chain = create_extraction_chain(
    llm, schema, encoder_or_encoder_class="json", validator=validator
)
print(chain.get_prompts()[0].format_prompt(text="[user input]").to_string())
ic(chain.invoke("stop the music now")["validated_data"])